import asyncio
from pysnmp.hlapi import *
import logging
from datetime import datetime
from pysnmp.hlapi.v3arch.asyncio import (
    SnmpEngine,
    CommunityData,
    UdpTransportTarget,
    ContextData,
    ObjectType,
    ObjectIdentity,
    set_cmd,
    get_cmd,
)

logger = logging.getLogger(__name__)


class SNMPClient:
    """SNMP客户端工具类，封装常用操作"""

    @staticmethod
    async def _create_transport_target(ip, timeout=2, retries=1):
        """创建传输目标的辅助方法"""
        return await UdpTransportTarget.create((ip, 161), timeout=timeout, retries=retries)

    @staticmethod
    async def _get_snmp_data(ip, community, oids, timeout=2, retries=1):
        """通用的SNMP获取数据方法"""
        transport_target = await SNMPClient._create_transport_target(ip, timeout, retries)
        snmp_engine = SnmpEngine()
        
        object_types = [ObjectType(ObjectIdentity(oid)) for oid in oids]
        result = await get_cmd(
            snmp_engine,
            CommunityData(community),
            transport_target,
            ContextData(),
            *object_types
        )
        
        return result

    @staticmethod
    def get_system_info(ip, community="public", timeout=2, retries=1):
        """
        获取设备基本信息
        :param ip: 设备IP
        :param community: SNMP团体名
        :return: 字典格式的系统信息 {name, cpu, status}
        """
        try:
            # 系统名称OID和CPU负载OID
            oids = ["1.3.6.1.2.1.1.5.0", "1.3.6.1.2.1.25.3.3.1.2"]
            
            result = asyncio.run(SNMPClient._get_snmp_data(ip, community, oids, timeout, retries))
            errorIndication, errorStatus, errorIndex, varBinds = result

            if errorIndication:
                raise ConnectionError(f"SNMP错误: {errorIndication}")
            elif errorStatus:
                raise ConnectionError(
                    f"SNMP状态错误: {errorStatus.prettyPrint()} at {errorIndex}"
                )

            return {
                "name": str(varBinds[0][1]),
                "cpu": int(varBinds[1][1]) if str(varBinds[1][1]).isdigit() else 0,
                "status": "在线",
                "timestamp": datetime.now().strftime("%Y-%m-%d %H:%M:%S"),
            }

        except Exception as e:
            logger.error(f"获取设备 {ip} 信息失败: {str(e)}")
            raise ConnectionError(f"SNMP请求异常: {str(e)}")

    @staticmethod
    async def _set_snmp_data(ip, community, oid, value, timeout=3, retries=1):
        """通用的SNMP设置数据方法"""
        transport_target = await SNMPClient._create_transport_target(ip, timeout, retries)
        snmp_engine = SnmpEngine()
        
        result = await set_cmd(
            snmp_engine,
            CommunityData(community),
            transport_target,
            ContextData(),
            ObjectType(ObjectIdentity(oid), value)
        )
        
        return result

    @staticmethod
    def send_flow_table(ip, community, flow_rules, timeout=3, retries=1):
        """
        下发流表到OpenDaylight
        :param ip: ODL控制器IP
        :param community: SNMP团体名
        :param flow_rules: 流表规则列表
        :return: bool 是否成功
        """
        try:
            ODL_FLOW_OID = "1.3.6.1.4.1.8072.9999.1.1"

            for rule in flow_rules:
                required_fields = {"table_id", "priority", "match", "actions"}
                if not all(field in rule for field in required_fields):
                    raise ValueError(f"流表规则缺少必要字段，需要: {required_fields}")

                safe_value = f"{rule['table_id']}|{rule['priority']}|{rule['match'].get('ipv4_destination','')}"

                result = asyncio.run(SNMPClient._set_snmp_data(
                    ip, community, ODL_FLOW_OID, safe_value, timeout, retries
                ))
                errorIndication, errorStatus, _, _ = result

                if errorIndication or errorStatus:
                    logger.error(f"下发流表失败: {errorIndication or errorStatus}")
                    return False

            logger.info(f"成功向 {ip} 下发 {len(flow_rules)} 条流表规则")
            return True

        except Exception as e:
            logger.error(f"流表下发异常: {str(e)}")
            return False

    @staticmethod
    def check_device_status(ip, community="public", timeout=2, retries=1):
        """快速检查设备是否在线"""
        try:
            result = asyncio.run(SNMPClient._get_snmp_data(
                ip, community, ["1.3.6.1.2.1.1.3.0"], timeout, retries
            ))
            errorIndication, _, _, _ = result
            return not errorIndication
        except:
            return False
